#[cfg(feature = "allocator_api")]
use std::alloc::Allocator;
use std::{mem::MaybeUninit, rc::Rc, sync::Arc};

use crate::*;

/// A trait for immutable buffers.
///
/// The `IoBuf` trait is implemented by buffer types that can be passed to
/// immutable completion-based IO operations, like writing its content to a
/// file. This trait will only take initialized bytes of a buffer into account.
pub trait IoBuf: 'static {
    /// Get the slice of initialized bytes.
    fn as_slice(&self) -> &[u8];

    /// Length of initialized bytes in the buffer.
    fn buf_len(&self) -> usize {
        self.as_slice().len()
    }

    /// Raw pointer to the buffer.
    fn buf_ptr(&self) -> *const u8 {
        self.as_slice().as_ptr()
    }

    /// Returns a view of the buffer with the specified range.
    ///
    /// This method is similar to Rust's slicing (`&buf[..]`), but takes
    /// ownership of the buffer.
    ///
    /// # Examples
    ///
    /// ```
    /// use compio_buf::IoBuf;
    ///
    /// let buf = b"hello world";
    /// assert_eq!(buf.slice(6..).as_slice(), b"world");
    /// ```
    fn slice(self, range: impl std::ops::RangeBounds<usize>) -> Slice<Self>
    where
        Self: Sized,
    {
        use std::ops::Bound;

        let begin = match range.start_bound() {
            Bound::Included(&n) => n,
            Bound::Excluded(&n) => n + 1,
            Bound::Unbounded => 0,
        };

        let end = match range.end_bound() {
            Bound::Included(&n) => Some(n.checked_add(1).expect("out of range")),
            Bound::Excluded(&n) => Some(n),
            Bound::Unbounded => None,
        };

        if let Some(end) = end {
            assert!(begin <= end);
        }

        Slice::new(self, begin, end)
    }
}

impl<B: IoBuf + ?Sized> IoBuf for &'static B {
    fn as_slice(&self) -> &[u8] {
        (**self).as_slice()
    }
}

impl<B: IoBuf + ?Sized> IoBuf for &'static mut B {
    fn as_slice(&self) -> &[u8] {
        (**self).as_slice()
    }
}

impl<B: IoBuf + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBuf
    for t_alloc!(Box, B, A)
{
    fn as_slice(&self) -> &[u8] {
        (**self).as_slice()
    }
}

impl<B: IoBuf + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBuf
    for t_alloc!(Rc, B, A)
{
    fn as_slice(&self) -> &[u8] {
        (**self).as_slice()
    }
}

impl IoBuf for [u8] {
    fn as_slice(&self) -> &[u8] {
        self
    }
}

impl<const N: usize> IoBuf for [u8; N] {
    fn as_slice(&self) -> &[u8] {
        self
    }
}

impl<#[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBuf for t_alloc!(Vec, u8, A) {
    fn as_slice(&self) -> &[u8] {
        Vec::as_slice(self)
    }
}

impl IoBuf for str {
    fn as_slice(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl IoBuf for String {
    fn as_slice(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl<B: IoBuf + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBuf
    for t_alloc!(Arc, B, A)
{
    fn as_slice(&self) -> &[u8] {
        (**self).as_slice()
    }
}

#[cfg(feature = "bytes")]
impl IoBuf for bytes::Bytes {
    fn as_slice(&self) -> &[u8] {
        (**self).as_slice()
    }
}

#[cfg(feature = "bytes")]
impl IoBuf for bytes::BytesMut {
    fn as_slice(&self) -> &[u8] {
        (**self).as_slice()
    }
}

#[cfg(feature = "read_buf")]
impl IoBuf for std::io::BorrowedBuf<'static> {
    fn as_slice(&self) -> &[u8] {
        self.filled()
    }
}

#[cfg(feature = "arrayvec")]
impl<const N: usize> IoBuf for arrayvec::ArrayVec<u8, N> {
    fn as_slice(&self) -> &[u8] {
        self
    }
}

#[cfg(feature = "smallvec")]
impl<const N: usize> IoBuf for smallvec::SmallVec<[u8; N]>
where
    [u8; N]: smallvec::Array<Item = u8>,
{
    fn as_slice(&self) -> &[u8] {
        self
    }
}

/// A trait for mutable buffers.
///
/// The `IoBufMut` trait is implemented by buffer types that can be passed to
/// mutable completion-based IO operations, like reading content from a file and
/// write to the buffer. This trait will take all space of a buffer into
/// account, including uninitialized bytes.
pub trait IoBufMut: IoBuf + SetBufInit {
    /// Get the full mutable slice of the buffer, including both initialized
    /// and uninitialized bytes.
    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>];

    /// Total capacity of the buffer, including both initialized and
    /// uninitialized bytes.
    fn buf_capacity(&mut self) -> usize {
        self.as_uninit().len()
    }

    /// Get the raw mutable pointer to the buffer.
    fn buf_mut_ptr(&mut self) -> *mut MaybeUninit<u8> {
        (*self).as_slice().as_ptr() as _
    }

    /// Get the mutable slice of initialized bytes. The content is the same as
    /// `as_slice`, but mutable.
    fn as_mut_slice(&mut self) -> &mut [u8] {
        let len = (*self).buf_len();
        let ptr = (*self).buf_mut_ptr();
        // SAFETY:
        // - lifetime of the returned slice is bounded by &mut self
        // - bytes within `len` are guaranteed to be initialized
        // - the pointer is derived from
        unsafe { std::slice::from_raw_parts_mut(ptr as *mut u8, len) }
    }

    /// Returns an [`Uninit`], which is a [`Slice`] that only exposes
    /// uninitialized bytes.
    ///
    /// It will always point to the uninitialized area of a [`IoBufMut`] even
    /// after reading in some bytes, which is done by [`SetBufInit`]. This
    /// is useful for writing data into buffer without overwriting any
    /// existing bytes.
    ///
    /// # Examples
    ///
    /// ```
    /// use compio_buf::{IoBuf, IoBufMut};
    ///
    /// let mut buf = Vec::from(b"hello world");
    /// buf.reserve_exact(10);
    /// let mut slice = buf.uninit();
    ///
    /// assert_eq!(slice.as_slice(), b"");
    /// assert_eq!(slice.buf_capacity(), 10);
    /// ```
    fn uninit(self) -> Uninit<Self>
    where
        Self: Sized,
    {
        Uninit::new(self)
    }

    /// Indicate whether the buffer has been filled (uninit portion is empty)
    fn is_filled(&mut self) -> bool {
        let len = (*self).as_slice().len();
        let cap = (*self).buf_capacity();
        len == cap
    }
}

impl<B: IoBufMut + ?Sized> IoBufMut for &'static mut B {
    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
        (**self).as_uninit()
    }
}

impl<B: IoBufMut + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBufMut
    for t_alloc!(Box, B, A)
{
    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
        (**self).as_uninit()
    }
}

impl<#[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBufMut for t_alloc!(Vec, u8, A) {
    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
        let cap = self.capacity();
        // SAFETY: Vec guarantees that the pointer is valid for `capacity` bytes
        unsafe { std::slice::from_raw_parts_mut(ptr, cap) }
    }
}

impl IoBufMut for [u8] {
    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
        let len = self.len();
        // SAFETY: slice is fully initialized, so treating it as MaybeUninit is safe
        unsafe { std::slice::from_raw_parts_mut(ptr, len) }
    }
}

impl<const N: usize> IoBufMut for [u8; N] {
    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
        // SAFETY: array is fully initialized, so treating it as MaybeUninit is safe
        unsafe { std::slice::from_raw_parts_mut(ptr, N) }
    }
}

#[cfg(feature = "bytes")]
impl IoBufMut for bytes::BytesMut {
    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
        let cap = self.capacity();
        // SAFETY: BytesMut guarantees that the pointer is valid for `capacity` bytes
        unsafe { std::slice::from_raw_parts_mut(ptr, cap) }
    }
}

#[cfg(feature = "read_buf")]
impl IoBufMut for std::io::BorrowedBuf<'static> {
    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
        let total_cap = self.capacity();

        // SAFETY: We reconstruct the full buffer from the filled portion pointer.
        // BorrowedBuf guarantees that the underlying buffer has capacity bytes.
        unsafe {
            let filled_ptr = self.filled().as_ptr() as *mut MaybeUninit<u8>;
            std::slice::from_raw_parts_mut(filled_ptr, total_cap)
        }
    }
}

#[cfg(feature = "arrayvec")]
impl<const N: usize> IoBufMut for arrayvec::ArrayVec<u8, N> {
    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
        // SAFETY: ArrayVec guarantees that the pointer is valid for N bytes
        unsafe { std::slice::from_raw_parts_mut(ptr, N) }
    }
}

#[cfg(feature = "smallvec")]
impl<const N: usize> IoBufMut for smallvec::SmallVec<[u8; N]>
where
    [u8; N]: smallvec::Array<Item = u8>,
{
    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
        let cap = self.capacity();
        // SAFETY: SmallVec guarantees that the pointer is valid for `capacity` bytes
        unsafe { std::slice::from_raw_parts_mut(ptr, cap) }
    }
}

/// A helper trait for `set_len` like methods.
pub trait SetBufInit {
    /// Set the buffer length. If `len` is less than the current length, this
    /// operation must be a no-op.
    ///
    /// # Safety
    ///
    /// * `len` must be less or equal than `as_uninit().len()`.
    /// * The bytes in the range `[buf_len(), len)` must be initialized.
    unsafe fn set_buf_init(&mut self, len: usize);
}

impl<B: SetBufInit + ?Sized> SetBufInit for &'static mut B {
    unsafe fn set_buf_init(&mut self, len: usize) {
        unsafe { (**self).set_buf_init(len) }
    }
}

impl<B: SetBufInit + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> SetBufInit
    for t_alloc!(Box, B, A)
{
    unsafe fn set_buf_init(&mut self, len: usize) {
        unsafe { (**self).set_buf_init(len) }
    }
}

impl<#[cfg(feature = "allocator_api")] A: Allocator + 'static> SetBufInit for t_alloc!(Vec, u8, A) {
    unsafe fn set_buf_init(&mut self, len: usize) {
        if (**self).buf_len() < len {
            unsafe { self.set_len(len) };
        }
    }
}

impl SetBufInit for [u8] {
    unsafe fn set_buf_init(&mut self, len: usize) {
        debug_assert!(len <= self.len());
    }
}

impl<const N: usize> SetBufInit for [u8; N] {
    unsafe fn set_buf_init(&mut self, len: usize) {
        debug_assert!(len <= N);
    }
}

#[cfg(feature = "bytes")]
impl SetBufInit for bytes::BytesMut {
    unsafe fn set_buf_init(&mut self, len: usize) {
        if (**self).buf_len() < len {
            unsafe { self.set_len(len) };
        }
    }
}

#[cfg(feature = "read_buf")]
impl SetBufInit for std::io::BorrowedBuf<'static> {
    #[allow(unsafe_op_in_unsafe_fn)]
    unsafe fn set_buf_init(&mut self, len: usize) {
        let current_len = (*self).buf_len();
        if current_len < len {
            // SAFETY: `len - current_len` range is initialized guaranteed by invariant of
            // `set_buf_init`
            unsafe { self.unfilled().advance(len - current_len) };
        }
    }
}

#[cfg(feature = "arrayvec")]
impl<const N: usize> SetBufInit for arrayvec::ArrayVec<u8, N> {
    unsafe fn set_buf_init(&mut self, len: usize) {
        if (**self).buf_len() < len {
            unsafe { self.set_len(len) };
        }
    }
}

#[cfg(feature = "smallvec")]
impl<const N: usize> SetBufInit for smallvec::SmallVec<[u8; N]>
where
    [u8; N]: smallvec::Array<Item = u8>,
{
    unsafe fn set_buf_init(&mut self, len: usize) {
        if (**self).buf_len() < len {
            unsafe { self.set_len(len) };
        }
    }
}

impl<T: IoBufMut> SetBufInit for [T] {
    unsafe fn set_buf_init(&mut self, len: usize) {
        unsafe { default_set_buf_init(self.iter_mut(), len) }
    }
}

impl<T: IoBufMut, const N: usize> SetBufInit for [T; N] {
    unsafe fn set_buf_init(&mut self, len: usize) {
        unsafe { default_set_buf_init(self.iter_mut(), len) }
    }
}

impl<T: IoBufMut, #[cfg(feature = "allocator_api")] A: Allocator + 'static> SetBufInit
    for t_alloc!(Vec, T, A)
{
    unsafe fn set_buf_init(&mut self, len: usize) {
        unsafe { default_set_buf_init(self.iter_mut(), len) }
    }
}

#[cfg(feature = "arrayvec")]
impl<T: IoBufMut, const N: usize> SetBufInit for arrayvec::ArrayVec<T, N> {
    unsafe fn set_buf_init(&mut self, len: usize) {
        unsafe { default_set_buf_init(self.iter_mut(), len) }
    }
}

#[cfg(feature = "smallvec")]
impl<T: IoBufMut, const N: usize> SetBufInit for smallvec::SmallVec<[T; N]>
where
    [T; N]: smallvec::Array<Item = T>,
{
    unsafe fn set_buf_init(&mut self, len: usize) {
        unsafe { default_set_buf_init(self.iter_mut(), len) }
    }
}

/// # Safety
/// * `len` should be less or equal than the sum of `buf_capacity()` of all
///   buffers.
/// * The bytes in the range `[buf_len(), new_len)` of each buffer must be
///   initialized
unsafe fn default_set_buf_init<'a, B: IoBufMut>(
    iter: impl IntoIterator<Item = &'a mut B>,
    mut len: usize,
) {
    let mut iter = iter.into_iter();
    while len > 0 {
        let Some(curr) = iter.next() else { return };
        let sub = (*curr).buf_capacity().min(len);
        unsafe { curr.set_buf_init(sub) };
        len -= sub;
    }
}
